import tensorflow as tf

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, warmup_steps=4000):
        super(CustomSchedule, self).__init__()

        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps ** -1.5)

        return tf.math.minimum(arg1, arg2)



learning_rate = CustomSchedule()

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                     epsilon=1e-9)
optimizer.iterations = tf.Variable(10)

x = tf.Variable(100000, dtype=tf.float32)

for i in range(10):
    with tf.GradientTape() as tape:
        y=10*x

    grads = tape.gradient(y, [x])
    optimizer.apply_gradients(zip(grads, [x]))

    checkpoint_path = "./checkpoints/train"
    ckpt = tf.train.Checkpoint(optimizer=optimizer)
    ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)
    ckpt.restore(ckpt_manager.latest_checkpoint)
for i in range(10):
	ckpt_save_path = ckpt_manager.save()
	print(ckpt_save_path)
